// YOLO8_SEG_RDK.hpp
#ifndef YOLOS_EDGEPLATFORM_YOLO8_SEG_RDK_HPP
#define YOLOS_EDGEPLATFORM_YOLO8_SEG_RDK_HPP

/**
 * @file YOLO8_SEG_RDK.hpp
 * @brief YOLO8 实例分割检测器（D-Robotics RDKx5 平台，BPU 硬件加速）
 *
 * 本头文件提供了针对 D-Robotics RDKx5 BPU（Brain Processing Unit）优化的
 * YOLO8 实例分割的 header-only 实现。
 *
 * 核心特性：
 * - 单头文件设计，易于集成
 * - BPU 硬件加速（通过 libDNN API）
 * - 支持 RGB (NCHW) / NV12 (YUV420SP) 输入格式自动检测
 * - DFL (Distribution Focal Loss) 后处理
 * - 高效的 mask 生成（支持 OpenMP 并行化）
 * - 可配置的预处理方式（Resize/LetterBox）
 * - RAII 资源管理
 *
 * 模型要求：
 * - 输入: NCHW-RGB/NV12, (1, 3, H, W), int8/int16 量化
 * - 输出: 10 个张量
 *   - 3 个尺度 (stride 8/16/32), 每个尺度 3 个输出:
 *     - cls[0,3,6]: (1, H//s, W//s, CLASSES_NUM) [NONE, float32]
 *     - bbox[1,4,7]: (1, H//s, W//s, 64) [SCALE, int32]
 *     - mce[2,5,8]: (1, H//s, W//s, 32) [SCALE, int32] - mask coefficients
 *   - proto[9]: (1, H//4, W//4, 32) [NONE/SCALE, float32/int16] - proto masks
 *
 * 使用示例：
 * @code
 * YOLO8SegDetector detector("model.bin", "coco.names");
 * cv::Mat image = cv::imread("test.jpg");
 * auto results = detector.detect(image, 0.25f, 0.7f, 0.5f);
 * for (auto& seg : results) {
 *     cv::rectangle(image, cv::Rect(seg.box.x, seg.box.y, seg.box.width, seg.box.height),
 *                   cv::Scalar(0,255,0), 2);
 *     // 使用 seg.mask 进行掩码处理
 * }
 * @endcode
 *
 * 作者: FANKYT
 * 日期: 2025
 */

#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
#include <fstream>
#include <cstring>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn/dnn.hpp>

// RDK BPU libDNN API
#include "dnn/hb_dnn.h"
#include "dnn/hb_dnn_ext.h"
#include "dnn/hb_sys.h"

#include "seg/BaseSegDetector.hpp"
#include "tools/ScopedTimer.hpp"
#include "tools/Common.hpp"

namespace yolos_edgeplatform {

/**
 * @brief YOLO8 实例分割检测器实现（RDK X5 平台）
 */
class YOLO8SegDetector : public BaseSegDetector {
public:
    /**
     * @brief 构造函数 - 加载 BPU 模型并初始化检测器
     *
     * @param modelPath .bin 模型文件路径
     * @param labelsPath 类别名称文件路径（每行一个类别）
     * @param numClasses 检测类别数量（默认: 80，COCO 数据集）
     * @param reg DFL 回归分箱数（默认: 16）
     * @param maskCoeffs mask 系数数量（默认: 32）
     */
    YOLO8SegDetector(const std::string &modelPath,
                     const std::string &labelsPath,
                     int numClasses = 80,
                     int reg = 16,
                     int maskCoeffs = 32)
        : numClasses_(numClasses), reg_(reg), maskCoeffs_(maskCoeffs),
          preprocessType_(PreprocessType::LETTERBOX) {

        ScopedTimer timer("Model Loading");

        // 加载类别名称
        classNames_ = loadClassNames(labelsPath);

        // 初始化 BPU 模型
        const char *modelFile = modelPath.c_str();
        int ret = hbDNNInitializeFromFiles(&packedDNNHandle_, &modelFile, 1);
        if (ret != 0) {
            throw std::runtime_error("[ERROR] hbDNNInitializeFromFiles failed with code: " +
                                     std::to_string(ret));
        }

        // 获取模型信息
        const char **modelNameList;
        int modelCount = 0;
        hbDNNGetModelNameList(&modelNameList, &modelCount, packedDNNHandle_);
        if (modelCount > 1) {
            std::cout << "[WARN] Multiple models in bin file, using first one" << std::endl;
        }

        const char *modelName = modelNameList[0];
        ret = hbDNNGetModelHandle(&dnnHandle_, packedDNNHandle_, modelName);
        if (ret != 0) {
            throw std::runtime_error("[ERROR] hbDNNGetModelHandle failed");
        }

        // 获取输入属性
        int32_t inputCount = 0;
        hbDNNGetInputCount(&inputCount, dnnHandle_);
        if (inputCount != 1) {
            throw std::runtime_error("[ERROR] Model should have exactly 1 input");
        }

        hbDNNGetInputTensorProperties(&inputProperties_, dnnHandle_, 0);

        // 自动检测输入张量类型并适配
        if (inputProperties_.tensorType == HB_DNN_IMG_TYPE_RGB) {
            inputImageType_ = InputImageType::RGB;
            std::cout << "[INFO] Input type: RGB (NCHW, int8 quantized)" << std::endl;
        } else if (inputProperties_.tensorType == HB_DNN_IMG_TYPE_NV12) {
            inputImageType_ = InputImageType::NV12;
            std::cout << "[INFO] Input type: NV12 (YUV420SP)" << std::endl;
        } else {
            throw std::runtime_error("[ERROR] Unsupported input tensor type");
        }

        // 验证输入张量布局
        if (inputProperties_.tensorLayout != HB_DNN_LAYOUT_NCHW) {
            throw std::runtime_error("[ERROR] Input tensor layout should be NCHW");
        }

        // 获取输入尺寸
        if (inputProperties_.validShape.numDimensions == 4) {
            inputH_ = inputProperties_.validShape.dimensionSize[2];
            inputW_ = inputProperties_.validShape.dimensionSize[3];
            std::cout << "[INFO] Input size: " << inputW_ << "x" << inputH_ << std::endl;
        } else {
            throw std::runtime_error("[ERROR] Invalid input tensor shape");
        }

        // 获取输出属性
        int32_t outputCount = 0;
        hbDNNGetOutputCount(&outputCount, dnnHandle_);
        if (outputCount != 10) {
            throw std::runtime_error("[ERROR] YOLO8-Seg model should have 10 outputs");
        }

        // 打印所有输出信息
        std::cout << "[INFO] Model outputs:" << std::endl;
        for (int i = 0; i < 10; i++) {
            hbDNNTensorProperties outputProps;
            hbDNNGetOutputTensorProperties(&outputProps, dnnHandle_, i);
            std::cout << "  output[" << i << "]: ("
                      << outputProps.validShape.dimensionSize[0] << ", "
                      << outputProps.validShape.dimensionSize[1] << ", "
                      << outputProps.validShape.dimensionSize[2] << ", "
                      << outputProps.validShape.dimensionSize[3] << "), ";
            if (outputProps.quantiType == SHIFT) std::cout << "SHIFT" << std::endl;
            else if (outputProps.quantiType == SCALE) std::cout << "SCALE" << std::endl;
            else if (outputProps.quantiType == NONE) std::cout << "NONE" << std::endl;
        }

        // 确定输出顺序
        determineOutputOrder();

        std::cout << "[INFO] YOLO8-Seg Detector initialized successfully" << std::endl;
        std::cout << "[INFO] Model: " << modelPath << std::endl;
        std::cout << "[INFO] Classes: " << numClasses_ << std::endl;
    }

    /**
     * @brief 析构函数 - 释放 BPU 资源
     */
    ~YOLO8SegDetector() override {
        if (packedDNNHandle_) {
            hbDNNRelease(packedDNNHandle_);
        }
    }

    /**
     * @brief 对图像执行实例分割检测
     *
     * @param image 输入图像（BGR 格式）
     * @param confThreshold 置信度阈值（默认 0.25）
     * @param nmsThreshold NMS IoU 阈值（默认 0.7）
     * @param maskThreshold 掩码二值化阈值（默认 0.5）
     * @return std::vector<Segmentation> 分割检测结果向量
     */
    std::vector<Segmentation> detect(const cv::Mat &image,
                                      float confThreshold = 0.25f,
                                      float nmsThreshold = 0.7f,
                                      float maskThreshold = 0.5f) override {
        ScopedTimer timer("Overall Detection");

        // 预处理
        cv::Mat processedImage;
        int xShift = 0, yShift = 0;
        float scale = 1.0f;
        preprocess(image, processedImage, xShift, yShift, scale);

        // 推理
        std::vector<hbDNNTensor> outputTensors = inference(processedImage);

        // 后处理
        std::vector<Segmentation> segmentations = postprocess(
            outputTensors, confThreshold, nmsThreshold, maskThreshold,
            xShift, yShift, scale, image.rows, image.cols);

        // 释放输出张量
        for (auto &tensor : outputTensors) {
            hbSysFreeMem(&tensor.sysMem[0]);
        }

        return segmentations;
    }

    cv::Size getInputSize() const override {
        return cv::Size(inputW_, inputH_);
    }

    int getNumClasses() const override {
        return numClasses_;
    }

    const std::vector<std::string>& getClassNames() const override {
        return classNames_;
    }

    void setPreprocessType(PreprocessType type) override {
        preprocessType_ = type;
    }

private:
    // BPU handles
    hbPackedDNNHandle_t packedDNNHandle_ = nullptr;
    hbDNNHandle_t dnnHandle_ = nullptr;
    hbDNNTensorProperties inputProperties_;

    // 模型参数
    int inputH_ = 640;
    int inputW_ = 640;
    int numClasses_ = 80;
    int reg_ = 16;           // DFL bins
    int maskCoeffs_ = 32;    // mask coefficients

    // 输出顺序映射（10个输出）
    int outputOrder_[10] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};

    // 配置
    PreprocessType preprocessType_;
    InputImageType inputImageType_;  ///< 记录模型的输入格式
    std::vector<std::string> classNames_;

    /**
     * @brief 自动确定输出张量顺序映射
     *
     * YOLO8-Seg 的 10 个输出:
     * - output[0,3,6]: cls (CLASSES_NUM), stride 8/16/32
     * - output[1,4,7]: bbox (64), stride 8/16/32
     * - output[2,5,8]: mce (32), stride 8/16/32
     * - output[9]: proto (32), H/4 x W/4
     *
     * 映射后的顺序:
     * - order[0-2]: cls_8, bbox_8, mce_8
     * - order[3-5]: cls_16, bbox_16, mce_16
     * - order[6-8]: cls_32, bbox_32, mce_32
     * - order[9]: proto
     */
    void determineOutputOrder() {
        int32_t H_4 = inputH_ / 4;
        int32_t H_8 = inputH_ / 8;
        int32_t H_16 = inputH_ / 16;
        int32_t H_32 = inputH_ / 32;
        int32_t W_4 = inputW_ / 4;
        int32_t W_8 = inputW_ / 8;
        int32_t W_16 = inputW_ / 16;
        int32_t W_32 = inputW_ / 32;

        // 期望的输出形状
        int32_t expectedShapes[10][3] = {
            {H_8, W_8, numClasses_},   // cls_8
            {H_8, W_8, 4 * reg_},      // bbox_8
            {H_8, W_8, maskCoeffs_},   // mce_8
            {H_16, W_16, numClasses_}, // cls_16
            {H_16, W_16, 4 * reg_},    // bbox_16
            {H_16, W_16, maskCoeffs_}, // mce_16
            {H_32, W_32, numClasses_}, // cls_32
            {H_32, W_32, 4 * reg_},    // bbox_32
            {H_32, W_32, maskCoeffs_}, // mce_32
            {H_4, W_4, maskCoeffs_}    // proto
        };

        // 匹配输出顺序
        for (int i = 0; i < 10; i++) {
            for (int j = 0; j < 10; j++) {
                hbDNNTensorProperties outputProps;
                hbDNNGetOutputTensorProperties(&outputProps, dnnHandle_, j);
                int32_t h = outputProps.validShape.dimensionSize[1];
                int32_t w = outputProps.validShape.dimensionSize[2];
                int32_t c = outputProps.validShape.dimensionSize[3];

                if (h == expectedShapes[i][0] &&
                    w == expectedShapes[i][1] &&
                    c == expectedShapes[i][2]) {
                    outputOrder_[i] = j;
                    break;
                }
            }
        }

        // 验证输出顺序
        int sum = 0;
        for (int i = 0; i < 10; i++) sum += outputOrder_[i];
        if (sum != 0+1+2+3+4+5+6+7+8+9) {
            std::cout << "[WARN] Output order check failed, using default order" << std::endl;
            for (int i = 0; i < 10; i++) outputOrder_[i] = i;
        } else {
            std::cout << "[INFO] Output order: {";
            for (int i = 0; i < 10; i++) {
                std::cout << outputOrder_[i] << (i < 9 ? ", " : "");
            }
            std::cout << "}" << std::endl;
        }
    }

    /**
     * @brief 预处理输入图像
     *
     * 根据模型输入类型自动选择预处理方法:
     * - RGB: letterBox/resize → BGR (保留,在 inference 转 RGB)
     * - NV12: letterBox/resize → BGR2YUV_I420 → I420转NV12
     */
    void preprocess(const cv::Mat &src, cv::Mat &dst,
                    int &xShift, int &yShift, float &scale) {
        ScopedTimer timer("Preprocessing");

        cv::Mat resizedImage;
        if (preprocessType_ == PreprocessType::LETTERBOX) {
            // LetterBox 预处理
            scale = std::min(static_cast<float>(inputH_) / src.rows,
                           static_cast<float>(inputW_) / src.cols);

            int newW = static_cast<int>(src.cols * scale);
            int newH = static_cast<int>(src.rows * scale);

            xShift = (inputW_ - newW) / 2;
            yShift = (inputH_ - newH) / 2;
            int xOther = inputW_ - newW - xShift;
            int yOther = inputH_ - newH - yShift;

            cv::resize(src, resizedImage, cv::Size(newW, newH));
            cv::copyMakeBorder(resizedImage, resizedImage, yShift, yOther, xShift, xOther,
                             cv::BORDER_CONSTANT, cv::Scalar(127, 127, 127));
        } else {
            // Resize 预处理
            cv::resize(src, resizedImage, cv::Size(inputW_, inputH_));
            scale = std::min(static_cast<float>(inputH_) / src.rows,
                           static_cast<float>(inputW_) / src.cols);
            xShift = 0;
            yShift = 0;
        }

        if (inputImageType_ == InputImageType::NV12) {
            // NV12 模式：BGR -> YUV I420 -> NV12
            cv::Mat yuvMat;
            cv::cvtColor(resizedImage, yuvMat, cv::COLOR_BGR2YUV_I420);
            uint8_t *yuv = yuvMat.ptr<uint8_t>();

            // I420 -> NV12 (YUV420SP)
            dst = cv::Mat(inputH_ * 3 / 2, inputW_, CV_8UC1);
            uint8_t *nv12Data = dst.ptr<uint8_t>();

            int ySize = inputH_ * inputW_;
            int uvHeight = inputH_ / 2;
            int uvWidth = inputW_ / 2;

            // 复制 Y 平面
            std::memcpy(nv12Data, yuv, ySize);

            // 交错复制 UV 平面
            uint8_t *nv12_uv = nv12Data + ySize;
            uint8_t *i420_u = yuv + ySize;
            uint8_t *i420_v = i420_u + uvHeight * uvWidth;

            for (int i = 0; i < uvHeight * uvWidth; i++) {
                nv12_uv[2 * i] = i420_u[i];       // U
                nv12_uv[2 * i + 1] = i420_v[i];   // V
            }
        } else {
            // RGB 模式：保持 BGR 格式，在 inference 中转换
            dst = resizedImage;
        }
    }

    /**
     * @brief 运行 BPU 推理
     *
     * 根据输入格式自动选择数据处理方式:
     * - RGB: BGR→RGB + NCHW + 量化(u8-128)
     * - NV12: 直接拷贝 NV12 数据
     */
    std::vector<hbDNNTensor> inference(const cv::Mat &image) {
        ScopedTimer timer("BPU Inference");

        // 分配输入张量
        hbDNNTensor inputTensor;
        inputTensor.properties = inputProperties_;

        if (inputImageType_ == InputImageType::RGB) {
            // RGB 模式：分配 3*H*W 内存
            hbSysAllocCachedMem(&inputTensor.sysMem[0], 3 * inputH_ * inputW_);

            // BGR -> RGB + NCHW + 量化 (u8-128 -> int8)
            uint8_t *srcData = image.ptr<uint8_t>();
            int8_t *dstData = reinterpret_cast<int8_t*>(inputTensor.sysMem[0].virAddr);

            for (int h = 0; h < inputH_; h++) {
                for (int w = 0; w < inputW_; w++) {
                    int srcIdx = h * inputW_ * 3 + w * 3;
                    int rIdx = (0 * inputH_ * inputW_) + h * inputW_ + w;
                    int gIdx = (1 * inputH_ * inputW_) + h * inputW_ + w;
                    int bIdx = (2 * inputH_ * inputW_) + h * inputW_ + w;

                    dstData[rIdx] = static_cast<int8_t>(srcData[srcIdx + 2] - 128); // R
                    dstData[gIdx] = static_cast<int8_t>(srcData[srcIdx + 1] - 128); // G
                    dstData[bIdx] = static_cast<int8_t>(srcData[srcIdx + 0] - 128); // B
                }
            }
        } else {
            // NV12 模式：分配 1.5*H*W 内存
            int inputSize = inputH_ * inputW_ * 3 / 2;
            hbSysAllocCachedMem(&inputTensor.sysMem[0], inputSize);

            // 直接复制 NV12 数据
            std::memcpy(inputTensor.sysMem[0].virAddr, image.ptr<uint8_t>(), inputSize);
        }

        hbSysFlushMem(&inputTensor.sysMem[0], HB_SYS_MEM_CACHE_CLEAN);

        // 分配输出张量
        std::vector<hbDNNTensor> outputTensors(10);
        for (int i = 0; i < 10; i++) {
            hbDNNGetOutputTensorProperties(&outputTensors[i].properties, dnnHandle_, i);
            int outAlignedSize = outputTensors[i].properties.alignedByteSize;
            hbSysAllocCachedMem(&outputTensors[i].sysMem[0], outAlignedSize);
        }

        // 运行推理
        hbDNNTaskHandle_t taskHandle = nullptr;
        hbDNNInferCtrlParam inferCtrlParam;
        HB_DNN_INITIALIZE_INFER_CTRL_PARAM(&inferCtrlParam);
        hbDNNInfer(&taskHandle, outputTensors.data(), &inputTensor, dnnHandle_, &inferCtrlParam);
        hbDNNWaitTaskDone(taskHandle, 0);
        hbDNNReleaseTask(taskHandle);

        // 释放输入张量
        hbSysFreeMem(&inputTensor.sysMem[0]);

        return outputTensors;
    }

    /**
     * @brief 后处理 BPU 输出生成分割结果
     *
     * 后处理流程:
     * 1. 解码 3 个尺度的检测输出 (bbox + mask coefficients)
     * 2. NMS 筛选
     * 3. 生成 mask (mask = sigmoid(mce @ proto.T))
     * 4. 上采样 mask 到原图尺寸
     * 5. 裁剪 mask 到 bbox 范围
     */
    std::vector<Segmentation> postprocess(
        std::vector<hbDNNTensor> &outputs,
        float confThreshold, float nmsThreshold, float maskThreshold,
        int xShift, int yShift, float scale,
        int origH, int origW) {

        ScopedTimer timer("Postprocessing");

        float confThreshRaw = -std::log(1.0f / confThreshold - 1.0f);

        // 临时存储解码结果
        std::vector<cv::Rect2d> decodedBboxes;
        std::vector<float> decodedScores;
        std::vector<int> decodedClasses;
        std::vector<std::vector<float>> decodedMces;

        // Step 1: 解析 Proto Mask
        int protoIdx = outputOrder_[9];
        hbSysFlushMem(&outputs[protoIdx].sysMem[0], HB_SYS_MEM_CACHE_INVALIDATE);

        int H_4 = inputH_ / 4;
        int W_4 = inputW_ / 4;
        std::vector<float> protoData(H_4 * W_4 * maskCoeffs_);

        // 根据量化类型获取 proto 数据
        if (outputs[protoIdx].properties.quantiType == NONE) {
            // float32 数据
            auto *protoFloat = reinterpret_cast<float*>(outputs[protoIdx].sysMem[0].virAddr);
            std::memcpy(protoData.data(), protoFloat, protoData.size() * sizeof(float));
        } else if (outputs[protoIdx].properties.quantiType == SCALE) {
            // int16 量化数据
            auto *protoRaw = reinterpret_cast<int16_t*>(outputs[protoIdx].sysMem[0].virAddr);
            float protoScale = outputs[protoIdx].properties.scale.scaleData[0];
            for (size_t i = 0; i < protoData.size(); i++) {
                protoData[i] = static_cast<float>(protoRaw[i]) * protoScale;
            }
        }

        // 创建 proto 矩阵 (H_4*W_4 x maskCoeffs_)
        cv::Mat protoMat(H_4 * W_4, maskCoeffs_, CV_32F, protoData.data());

        // Step 2: 处理 3 个检测头
        int strides[3] = {8, 16, 32};
        int heights[3] = {inputH_ / 8, inputH_ / 16, inputH_ / 32};
        int widths[3] = {inputW_ / 8, inputW_ / 16, inputW_ / 32};

        for (int headIdx = 0; headIdx < 3; headIdx++) {
            int clsIdx = outputOrder_[headIdx * 3];
            int bboxIdx = outputOrder_[headIdx * 3 + 1];
            int mceIdx = outputOrder_[headIdx * 3 + 2];

            // 刷新内存
            hbSysFlushMem(&outputs[clsIdx].sysMem[0], HB_SYS_MEM_CACHE_INVALIDATE);
            hbSysFlushMem(&outputs[bboxIdx].sysMem[0], HB_SYS_MEM_CACHE_INVALIDATE);
            hbSysFlushMem(&outputs[mceIdx].sysMem[0], HB_SYS_MEM_CACHE_INVALIDATE);

            // 获取数据指针
            auto *clsRaw = reinterpret_cast<float*>(outputs[clsIdx].sysMem[0].virAddr);
            auto *bboxRaw = reinterpret_cast<int32_t*>(outputs[bboxIdx].sysMem[0].virAddr);
            auto *bboxScale = reinterpret_cast<float*>(outputs[bboxIdx].properties.scale.scaleData);
            auto *mceRaw = reinterpret_cast<int32_t*>(outputs[mceIdx].sysMem[0].virAddr);
            auto *mceScale = reinterpret_cast<float*>(outputs[mceIdx].properties.scale.scaleData);

            int H = heights[headIdx];
            int W = widths[headIdx];
            int stride = strides[headIdx];

            for (int h = 0; h < H; h++) {
                for (int w = 0; w < W; w++) {
                    int offset = h * W + w;
                    float *curCls = clsRaw + offset * numClasses_;
                    int32_t *curBbox = bboxRaw + offset * (4 * reg_);
                    int32_t *curMce = mceRaw + offset * maskCoeffs_;

                    // 找到最大置信度的类别
                    int bestClsId = 0;
                    for (int c = 1; c < numClasses_; c++) {
                        if (curCls[c] > curCls[bestClsId]) {
                            bestClsId = c;
                        }
                    }

                    // 提前过滤低置信度
                    if (curCls[bestClsId] < confThreshRaw) {
                        continue;
                    }

                    // Sigmoid 激活计算置信度
                    float conf = 1.0f / (1.0f + std::exp(-curCls[bestClsId]));

                    // DFL 解码边界框
                    float ltrb[4] = {0.0f, 0.0f, 0.0f, 0.0f};
                    for (int i = 0; i < 4; i++) {
                        float sum = 0.0f;
                        for (int j = 0; j < reg_; j++) {
                            int idx = reg_ * i + j;
                            float dfl = std::exp(static_cast<float>(curBbox[idx]) * bboxScale[idx]);
                            ltrb[i] += dfl * j;
                            sum += dfl;
                        }
                        ltrb[i] /= sum;
                    }

                    // 剔除不合格的框
                    if (ltrb[2] + ltrb[0] <= 0 || ltrb[3] + ltrb[1] <= 0) {
                        continue;
                    }

                    // 距离转换为边界框坐标
                    float x1 = (w + 0.5f - ltrb[0]) * stride;
                    float y1 = (h + 0.5f - ltrb[1]) * stride;
                    float x2 = (w + 0.5f + ltrb[2]) * stride;
                    float y2 = (h + 0.5f + ltrb[3]) * stride;

                    // 边界检查
                    if (x1 >= 0 && y1 >= 0 && x2 > x1 && y2 > y1 &&
                        x2 <= inputW_ && y2 <= inputH_) {
                        // 反量化 mask coefficients
                        std::vector<float> maskCoeffs(maskCoeffs_);
                        for (int i = 0; i < maskCoeffs_; i++) {
                            maskCoeffs[i] = static_cast<float>(curMce[i]) * mceScale[i];
                        }

                        decodedBboxes.push_back(cv::Rect2d(x1, y1, x2 - x1, y2 - y1));
                        decodedScores.push_back(conf);
                        decodedClasses.push_back(bestClsId);
                        decodedMces.push_back(maskCoeffs);
                    }
                }
            }
        }

        // Step 3: NMS 筛选
        std::vector<int> nmsIndices;
        if (!decodedBboxes.empty()) {
            std::vector<cv::Rect> nmsBboxes;
            for (const auto &box : decodedBboxes) {
                nmsBboxes.push_back(cv::Rect(
                    static_cast<int>(box.x),
                    static_cast<int>(box.y),
                    static_cast<int>(box.width),
                    static_cast<int>(box.height)
                ));
            }
            cv::dnn::NMSBoxes(nmsBboxes, decodedScores, confThreshold, nmsThreshold, nmsIndices);
        }

        // Step 4: 生成 masks 并构建最终结果
        std::vector<Segmentation> finalSegmentations;

        for (int idx : nmsIndices) {
            Segmentation seg;
            cv::Rect2d bbox = decodedBboxes[idx];

            // 生成 mask
            cv::Mat mce(1, maskCoeffs_, CV_32F, decodedMces[idx].data());
            cv::Mat maskFlat = protoMat * mce.t();  // (H_4*W_4, 1)
            cv::Mat maskLowRes = maskFlat.reshape(1, H_4);  // (H_4, W_4)

            // Sigmoid 激活
            cv::Mat sigmoidMask;
            cv::exp(-maskLowRes, sigmoidMask);
            sigmoidMask = 1.0 / (1.0 + sigmoidMask);

            // 上采样到模型输入尺寸
            cv::Mat resizedMask;
            cv::resize(sigmoidMask, resizedMask, cv::Size(inputW_, inputH_), 0, 0, cv::INTER_LINEAR);

            // 二值化
            cv::Mat binaryMask;
            cv::threshold(resizedMask, binaryMask, maskThreshold, 1.0, cv::THRESH_BINARY);
            binaryMask.convertTo(binaryMask, CV_8U, 255);

            // 裁剪到 bbox 区域
            int x1 = std::max(0, static_cast<int>(bbox.x));
            int y1 = std::max(0, static_cast<int>(bbox.y));
            int x2 = std::min(inputW_, static_cast<int>(bbox.x + bbox.width));
            int y2 = std::min(inputH_, static_cast<int>(bbox.y + bbox.height));

            if (x2 > x1 && y2 > y1) {
                cv::Rect roiRect(x1, y1, x2 - x1, y2 - y1);
                cv::Mat croppedMask = binaryMask(roiRect).clone();

                // 映射到原图坐标
                float origX1 = (x1 - xShift) / scale;
                float origY1 = (y1 - yShift) / scale;
                float origX2 = (x2 - xShift) / scale;
                float origY2 = (y2 - yShift) / scale;

                // 限制在原图范围内
                origX1 = std::max(0.0f, origX1);
                origY1 = std::max(0.0f, origY1);
                origX2 = std::min(static_cast<float>(origW), origX2);
                origY2 = std::min(static_cast<float>(origH), origY2);

                // 上采样 mask 到原图尺寸
                int origMaskW = static_cast<int>(origX2 - origX1);
                int origMaskH = static_cast<int>(origY2 - origY1);
                if (origMaskW > 0 && origMaskH > 0) {
                    cv::Mat finalMask = cv::Mat::zeros(origH, origW, CV_8U);
                    cv::Mat resizedCroppedMask;
                    cv::resize(croppedMask, resizedCroppedMask,
                              cv::Size(origMaskW, origMaskH), 0, 0, cv::INTER_LINEAR);

                    cv::Rect origRoi(static_cast<int>(origX1), static_cast<int>(origY1),
                                    origMaskW, origMaskH);
                    resizedCroppedMask.copyTo(finalMask(origRoi));

                    seg.box.x = origX1;
                    seg.box.y = origY1;
                    seg.box.width = origX2 - origX1;
                    seg.box.height = origY2 - origY1;
                    seg.conf = decodedScores[idx];
                    seg.classId = decodedClasses[idx];
                    seg.mask = finalMask;

                    finalSegmentations.push_back(seg);
                }
            }
        }

        return finalSegmentations;
    }
};

} // namespace yolos_edgeplatform

#endif // YOLOS_EDGEPLATFORM_YOLO8_SEG_RDK_HPP
